import os
import torch
import numpy as np

def trace(x):
    result = []
    for i in range(x.shape[0]):
        result.append(torch.trace(x[i]))
    return torch.stack(result)

def gen_golden_data_simple():
    input_x = np.random.uniform(-10, 10, [8, 10, 10]).astype(np.float16)
    golden = trace(torch.from_numpy(input_x.astype(np.float32))).numpy().astype(np.float16)
    os.system("mkdir -p input")
    os.system("mkdir -p output")

    input_x.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")

    print(f"Input shape: {input_x.shape}")
    print(f"Golden shape: {golden.shape}")
    print(f"Golden values: {golden}")

if __name__ == "__main__":
    gen_golden_data_simple()